from src.document_parser.parser_factory import DocumentParserFactory
from typing import Tuple, List, Dict
import os
import fitz
import logging
from operator import itemgetter

class DocumentParser:
    """文档解析器"""
    
    def __init__(self, config: dict):
        self.config = config
        self.header_height_ratio = 0.1  # 页眉区域高度比例
        self.footer_height_ratio = 0.1  # 页脚区域高度比例
        # 添加水印特征
        self.watermark_features = {
            'min_width_ratio': 0.5,    # 水印通常横跨页面一半以上
            'min_height_ratio': 0.3,    # 水印通常纵跨页面三分之一以上
            'transparency_threshold': 0.3,  # 水印通常是半透明的
            'patterns': ['厦门攸信信息技术有限公司','Ximen UMS Information Technology Co..LID.']  # 水印常见文字
        }
        self.ignored_patterns = [        # 需要忽略的文本模式
            r'\d+/\d+',                  # 页码格式 如 1/10
            r'第\s*\d+\s*页',            # 页码格式 如 第1页
            r'©.*版权所有',              # 版权信息
            r'机密文件',                 # 文件标记
            r'内部文件',
            r'攸信',                     # 公司名称
            r'uMOM'                      # 产品名称
        ]
        
    def _is_header_or_footer(self, block, page_height, page_width):
        """判断文本块是否是页眉页脚
        
        Args:
            block: 文本块信息
            page_height: 页面高度
            page_width: 页面宽度
            
        Returns:
            bool: 是否是页眉页脚
        """
        x0, y0, x1, y1 = block[:4]
        text = block[4]
        
        # 计算页眉页脚区域
        header_zone = page_height * self.header_height_ratio
        footer_zone = page_height * (1 - self.footer_height_ratio)
        
        # 位置判断
        is_in_header_zone = y0 < header_zone
        is_in_footer_zone = y1 > footer_zone
        
        # 内容判断
        import re
        is_ignored_content = any(re.search(pattern, text, re.IGNORECASE) for pattern in self.ignored_patterns)
        
        # 其他特征判断
        is_single_line = (y1 - y0) < 20  # 单行文本
        is_wide_text = (x1 - x0) > page_width * 0.5  # 跨页文本
        
        return (is_in_header_zone or is_in_footer_zone) and (is_ignored_content or is_single_line or is_wide_text)
        
    def _is_watermark(self, pix: fitz.Pixmap, rect: fitz.Rect, page_width: float, page_height: float) -> bool:
        """判断图片是否是水印"""
        try:
            # 1. 基于固定尺寸的判断
            WATERMARK_SIZE_THRESHOLD = 424
            if (rect.width >= WATERMARK_SIZE_THRESHOLD and 
                rect.height >= WATERMARK_SIZE_THRESHOLD and
                abs(rect.width - rect.height) < 50):  # 允许宽高有少许差异
                logging.info(f"基于尺寸判定为水印: width={rect.width}, height={rect.height}")
                return True
            
            # 2. 保留原有的颜色和透明度判断作为补充
            # 1. 基本尺寸检查
            width_ratio = rect.width / page_width
            height_ratio = rect.height / page_height
            
            # 2. 检查是否覆盖大部分页面
            if width_ratio < 0.5 or height_ratio < 0.5:  # 如果图片太小，不是水印
                return False
            
            # 3. 检查颜色特征
            if not (pix.alpha and pix.n >= 4):  # 必须有alpha通道
                return False
            
            samples = pix.samples
            pixel_count = pix.width * pix.height
            
            # 计算颜色统计
            color_stats = {
                'light_gray_pixels': 0,  # 浅灰色像素数
                'transparent_pixels': 0,  # 透明像素数
                'total_pixels': pixel_count
            }
            
            # 遍历像素
            for i in range(0, len(samples), 4):
                r, g, b, a = samples[i:i+4]
                
                # 检查是否是浅灰色 (R≈G≈B，且值较高)
                is_light_gray = (abs(r - g) < 20 and 
                               abs(g - b) < 20 and 
                               abs(b - r) < 20 and 
                               180 < r < 240)  # 浅灰色范围
                               
                # 检查透明度
                is_transparent = a < 128  # 半透明
                
                if is_light_gray:
                    color_stats['light_gray_pixels'] += 1
                if is_transparent:
                    color_stats['transparent_pixels'] += 1
            
            # 计算比例
            gray_ratio = color_stats['light_gray_pixels'] / color_stats['total_pixels']
            transparent_ratio = color_stats['transparent_pixels'] / color_stats['total_pixels']
            
            # 水印判定条件：
            # 1. 大尺寸（覆盖大部分页面）
            # 2. 主要是浅灰色
            # 3. 大部分是半透明的
            if (width_ratio > 0.5 and height_ratio > 0.5 and  # 大尺寸
                gray_ratio > 0.6 and                          # 主要是浅灰色
                transparent_ratio > 0.5):                     # 大部分半透明
                logging.info(f"检测到水印: size_ratio=({width_ratio:.2f}, {height_ratio:.2f}), "
                            f"gray_ratio={gray_ratio:.2f}, transparent_ratio={transparent_ratio:.2f}")
                return True
            
            return False
            
        except Exception as e:
            logging.warning(f"水印检查失败: {str(e)}")
            return False

    def parse_to_markdown(self, file_path: str, image_handler, source_name: str) -> str:
        """解析文档为Markdown格式"""
        try:
            doc = fitz.open(file_path)
            all_content = []
            
            # 遍历每一页
            for page_num in range(doc.page_count):
                page = doc[page_num]
                page_content = []
                
                # 获取页面尺寸
                page_width = page.rect.width
                page_height = page.rect.height
                
                # 1. 提取文本块，保留位置信息
                text_blocks = page.get_text("blocks")
                for block in text_blocks:
                    if block[6] == 0:  # 文本块
                        if not self._is_header_or_footer(block, page_height, page_width):
                            page_content.append({
                                'type': 'text',
                                'content': block[4],
                                'bbox': block[:4],
                                'y_pos': block[1]
                            })
                
                # 2. 提取图片，记录位置信息
                try:
                    # 获取页面上所有图片
                    page_images = []
                    img_list = page.get_images()
                    
                    for img_index, img in enumerate(img_list, start=1):
                        try:
                            xref = img[0]  # 获取图片的xref
                            if xref <= 0:
                                continue
                                
                            base_image = doc.extract_image(xref)
                            if not base_image:
                                continue
                                
                            # 获取图片在页面上的位置
                            pix = fitz.Pixmap(doc, xref)
                            image_rects = page.get_image_rects(img)  # 获取所有图片实例的位置
                            
                            if image_rects:  # 如果找到了图片位置
                                for rect in image_rects:
                                    # 过滤页眉页脚区域的图片
                                    if not (rect.y0 < page_height * self.header_height_ratio or 
                                           rect.y1 > page_height * (1 - self.footer_height_ratio)):
                                        # 检查图片尺寸
                                        width = rect.width
                                        height = rect.height
                                        if width > 50 and height > 50:  # 基本尺寸检查
                                            # 水印检查
                                            if not self._is_watermark(pix, rect, page_width, page_height):
                                                image_data = {
                                                    'image_data': base_image['image'],
                                                    'format': base_image['ext'].lower(),
                                                    'page': page_num + 1,
                                                    'index': len(page_images) + 1
                                                }
                                                
                                                page_content.append({
                                                    'type': 'image',
                                                    'data': image_data,
                                                    'bbox': rect,
                                                    'y_pos': rect.y0
                                                })
                                                page_images.append(image_data)
                                                break  # 只使用第一个实例的位置
                                            
                        except Exception as e:
                            logging.error(f"提取单个图片失败: page={page_num+1}, img_index={img_index}, error={str(e)}")
                            continue
                            
                except Exception as e:
                    logging.error(f"提取页面图片失败: page={page_num+1}, error={str(e)}")
                    continue
                
                # 3. 按y坐标排序页面内容
                if page_content:
                    page_content.sort(key=itemgetter('y_pos'))
                    
                    # 4. 生成页面的markdown内容
                    content_text = [f""]
                    images = []
                    
                    last_y_pos = -1
                    for item in page_content:
                        # 添加适当的空行
                        if last_y_pos >= 0 and item['y_pos'] - last_y_pos > 50:
                            content_text.append("\n")
                        
                        if item['type'] == 'text':
                            text = item['content'].strip()
                            if text:  # 只添加非空文本
                                content_text.append(text + "\n\n")
                        else:  # image
                            images.append(item['data'])
                            content_text.append(f"\n![](@page_{item['data']['page']}_image_{item['data']['index']})\n\n")
                        
                        last_y_pos = item['y_pos']
                    
                    # 5. 添加到总内容中
                    all_content.extend(images)
                    all_content.append(''.join(content_text))
            
            # 处理图片并获取URL
            image_urls = image_handler.process_images([img for img in all_content if isinstance(img, dict)], source_name)
            
            # 合并所有文本内容
            content = ''.join([c for c in all_content if isinstance(c, str)])
            
            # 替换图片引用为URL
            if image_urls:
                content = image_handler.replace_image_references(content, image_urls)
            
            # 保存为txt文件
            output_txt = f"markdown_{os.path.splitext(os.path.basename(file_path))[0]}.txt"
            with open(output_txt, 'w', encoding='utf-8') as f:
                f.write(content)
            
            return content
            
        except Exception as e:
            logging.error(f"解析文档失败: {str(e)}")
            raise e 